import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch

def draw_amha_improved():
    plt.rcParams.update({
        'font.sans-serif': 'Arial',
        'font.size': 9,
        'axes.linewidth': 1.2
    })
    
    fig = plt.figure(figsize=(14, 8))
    ax = fig.add_subplot(111)
    ax.axis('off')
    
    # 专业配色方案
    colors = {
        'input': '#2F5597',
        'downsample': '#8FAADC',
        'axial': '#C55A11',
        'channel': '#548235',
        'upsample': '#A9D08E',
        'fusion': '#7030A0',
        'arrow': '#7F7F7F'
    }
    
    # 模块布局参数
    module_config = {
        'input': (0.1, 0.8, 0.12, 0.1),
        'downsample': (0.3, 0.75, 0.15, 0.06),
        'axial_conv': (0.3, 0.6, 0.15, 0.12),
        'channel_attn': (0.3, 0.4, 0.15, 0.12),
        'upsample': (0.5, 0.6, 0.15, 0.06),
        'fusion': (0.7, 0.8, 0.12, 0.1)
    }
    
    # 绘制模块的增强函数
    def draw_module(x, y, w, h, label, color, text_offset=0.03, 
                   text_align='center', box_style='round'):
        box = FancyBboxPatch((x, y), w, h, 
                           boxstyle=f"{box_style},pad=0.02", 
                           facecolor=color, edgecolor='k', linewidth=1.2)
        ax.add_patch(box)
        
        # 多行文本对齐
        lines = label.split('\n')
        for i, line in enumerate(reversed(lines)):
            y_text = y + h/2 - i*text_offset
            plt.text(x + w/2, y_text, line, 
                    ha=text_align, va='center', fontsize=9)
        return box
    
    # 绘制所有模块
    modules = {}
    for key in module_config:
        params = module_config[key]
        modules[key] = draw_module(*params, 
                                  label=key.replace('_', ' ').title(),
                                  color=colors[key.split('_')[0]])
    
    # 修正后的连接线函数
    def connect_blocks(src, dst, style='straight', text=None):
        # 计算源模块中心坐标
        src_box = modules[src]
        x1 = src_box.get_x() + src_box.get_width()/2
        y1 = src_box.get_y() + src_box.get_height()/2
        
        # 计算目标模块中心坐标
        dst_box = modules[dst]
        x2 = dst_box.get_x() + dst_box.get_width()/2
        y2 = dst_box.get_y() + dst_box.get_height()/2
        
        arrow_style = dict(arrowstyle="->", 
                          color=colors['arrow'],
                          linewidth=1.2,
                          connectionstyle=f"arc3,rad={0.2 if style=='curved' else 0}",
                          shrinkA=8, shrinkB=8)
        
        ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
                   arrowprops=arrow_style)
        
        if text:
            plt.text((x1+x2)/2, (y1+y2)/2+0.02, text, 
                    ha='center', va='bottom', fontsize=8,
                    backgroundcolor='white')

    # 连接模块
    connect_blocks('input', 'downsample', text="C×H×W → C/2×H/2×W/2")
    connect_blocks('downsample', 'axial_conv')
    connect_blocks('axial_conv', 'upsample')
    connect_blocks('upsample', 'fusion', style='curved')
    connect_blocks('input', 'fusion', text="Skip Connection")
    
    # 添加数学标注
    plt.text(0.3, 0.55, r'$\otimes$', fontsize=20, ha='center', va='center')
    plt.text(0.3, 0.35, r'$Attention_{\mathrm{channel}}$', 
            fontsize=10, ha='center', backgroundcolor='white')
    
    # 添加比例尺
    ax.add_patch(plt.Rectangle((0.1, 0.1), 0.2, 0.02, color='#555555'))
    plt.text(0.2, 0.12, "Scale: 1 unit = 2 Conv Blocks", 
            ha='center', va='bottom', fontsize=8)
    
    plt.title("AMHA Module Architecture (Improved Visualization)", 
            fontsize=14, pad=20)
    plt.tight_layout()
    plt.savefig('AMHA.svg', bbox_inches='tight')
    plt.show()

draw_amha_improved()